package simplytypedinference


/**
 * Following typechecking as in Types and Programming Languages by Benjamin Pierce, chap 9-10-11
 *
 */

class TypeExc(s: String) extends Exception(s) 


object varnamegen extends Function1[String, String] {
  var counter: int = 0
  def apply(s:String):String = {
    counter += 1
    s+"%"+counter
  }
}

object LTerm {
  
  import inference._
  import inference.Constraint
  type Constraint = (LType, LType)

  
  def getTyConstraints(t: LTerm) : (List[Constraint], LType)= 
    getTyConstraints(new Context[VarBinding](Nil), t)
    
  def getTyConstraints(ctx:Context[VarBinding], t: LTerm): (List[Constraint], LType) = t match {
    case Lit(x)     => (Nil, TyInt)
    case Tru()      => (Nil, TyBool)
    case Fls()      => (Nil, TyBool) 
    case Unit()     => (Nil, TyUnit)
    
    case V(n)       => {
      try {(Nil, ctx(n).ltype)}
      catch {
        case e: Exception => throw new TypeExc("unknown var: "+n)
      }
    }
    case L(x,ty,t) => {
      val ty1 =  TyVar(varnamegen(x))
      val ctx2 : Context[VarBinding] = ctx + (x,new VarBinding(ty1))
      val (cs, ty2) = getTyConstraints(ctx2, t)
      (cs, TyArr(ty1, ty2))
    }
    case A(t1,t2)  => {
      val (cs1,ty1)   = getTyConstraints(ctx,t1)
      val (cs2,ty2)   = getTyConstraints(ctx,t2)
      val tyn =  TyVar(varnamegen("a"))
      val newcon :Constraint =  ty1 -> TyArr(ty2, tyn)
      val cons:List[Constraint] = newcon :: (cs2 ::: cs1)
      
      (cons, tyn)
    }
    
    case Let(Name(str),t1,t2) => {
      val (cs1, typeInsert) = getTyConstraints(ctx,t1)
      val ctx2 = ctx + (str, new VarBinding(typeInsert))
      val (cs2, tyt2) = getTyConstraints(ctx2, t2)
      val cons:List[Constraint] = cs2 ::: cs1
      
      (cons, tyt2)
    }
    case If(cond,pos,neg) => {
      val (cs_c, ty_c) = getTyConstraints(ctx,cond)
      val (cs_p, ty_p) = getTyConstraints(ctx, pos)
      val (cs_n, ty_n) = getTyConstraints(ctx, neg)
      val c_condbool = ty_c -> TyBool
      val c_branchequal = ty_p -> ty_n
      
      val cons:List[Constraint] = List(c_condbool, c_branchequal) ::: cs_c ::: cs_p ::: cs_n
      
      (cons, ty_p)
    }
  }
 
  def typeOf(t: LTerm):LType = {
    val (cons, ty) = getTyConstraints(t)
    val subst = unify(cons)
    
    var fullpart = ty
    println("   :: "+ty)
    println("subst:"+subst)
    subst.foreach(x => {fullpart = x.subst(fullpart)})
    fullpart
  
  }
  /*
  private def typeOf(t: LTerm, ctx: Context[VarBinding]) : LType= t match {
    case Lit(x)     => TyInt
    case Tru()      => TyBool
    case Fls()      => TyBool 
    case Unit()     => TyUnit
    case If(cond,pos,neg) => {
      if (!(TyBool equals typeOf(cond,ctx) )) throw new TypeExc("conditional not boolean")
      val typepos = typeOf(pos,ctx);val typeneg = typeOf(neg,ctx)
      if (!( typepos equals typeneg)) throw new TypeExc("branches of if differ in type: "+typepos+" vs "+typeneg)
      typepos
    }
    case Let(Name(str),t1,t2) => {
      val typeInsert = typeOf(t1, ctx)
      val ctx2 = ctx + (str, new VarBinding(typeInsert))
      typeOf(t2, ctx2)
    }
    case V(n)       => {
      try {ctx(n).ltype}
      catch {
        case e: Exception => null
      }
    }
    case L(x,ty,t) => {
      val ctx2 : Context[VarBinding] = ctx + (x,new VarBinding(ty))
      TyArr(ty, LTerm.typeOf(t,ctx2))
    }
    case A(t1,t2)  => {
      val ty1 = LTerm.typeOf(t1,ctx)
      val ty2 = LTerm.typeOf(t2,ctx)
      ty1 match {
        case TyArr(ty1D, ty1G) => {  
          if (ty1D equals ty2) ty1G
          else throw new TypeExc("receive "+ (if(ty2==null) "untyped" else ty2) 
              +" but expect "+ (if(ty1D==null)"" else ty1D))
        }
        case _                 => {
          throw new TypeExc("applying to non function type")
        }
      }
    }
  }
  
  */
  
  def eval(t: LTerm) : LTerm= {
    try {
    val t_bis = eval1 (t)
    eval (t_bis) // Q: is this tail-recursive?
    } 
    catch { case e: NoReductionPossibleExc => t }
  }
    
    
  case class NoReductionPossibleExc extends Exception
    
  def eval1(t: LTerm): LTerm = t match {
    case If(c,p,n) if (isValue(c))            =>
         if (c equals Tru) p else n
    case If(c,p,n)                            => If(eval1(c),p,n)
    case A( L(s,ty,t12), v2) if (isValue(v2)) => termSubstTop(v2,t12)
    case A(t1, t2)                            =>
         if (isValue(t1) ) 
              A(t1 ,eval1(t2))
        else 
              A(eval1(t1), t2)
    case Let(Name(n),v1,t2) if (isValue(v1))  => termSubst(0,v1,t2)
    case Let(Name(n),t1,t2)                   => Let(Name(n),eval1(t1),t2)
    case _                                    =>
         throw NoReductionPossibleExc
  }
  
  
  //types are actually irrelevant during evaluation, could use null everywhere
  def isValue(t: LTerm) = t match {
    case L(_,_,_) => true
    case Lit(v)   => true
    case Tru()    => true
    case Fls()    => true
    case Unit()   => true
    case _        => false
  }
  
  def termSubstTop(v: LTerm, t: LTerm): LTerm = {
    termShift(-1, 
              termSubst(0, 
                        termShift(1,v),
                        t ) )
  }
  def termShift(d: Int, t: LTerm) = {
  /*  def walk(c: Int,t: LTerm) :LTerm = t match {
      case V(nb)         => if (nb >= c) V(nb+d) else V(nb)
      case L(x,ty,t1)    => L(x,ty, walk(c+1, t1))
      case A(t1,t2)      => A(walk(c,t1),walk(c,t2))
      case Tru()         => Tru
      case Fls()         => Fls
      case Lit(v)        => Lit(v)
      case If(cond,p,n)  => If(walk(c,cond),walk(c,p),walk(c,n))  
    }
    walk(0,t)
  */  
    termMap((c:Int,x:V) => if (x.n >=c) V(x.n+d) else V(x.n) ,d,t)
  }
  
  def termSubst(j: Int, s: LTerm,t: LTerm) = {
  /*  def walk(c: Int, t: LTerm):LTerm = t match {
      case V(nb)         => if (nb == j+c) termShift(c,s) else V(nb)
      case L(x,ty,t1)    => L(x,ty, walk(c+1, t1))
      case A(t1,t2)      => A(walk(c,t1),walk(c,t2))
      case Tru()         => Tru
      case Fls()         => Fls
      case Lit(v)        => Lit(v)
      case If(cond,p,n)  => If(walk(c,cond),walk(c,p),walk(c,n))
    }
    walk(0,t)
  */
    termMap((c:Int,x:V) => if (x.n == j+c) termShift(c,s) else V(x.n),j,t)
  }
  def termMap(onvar:(Int,V)=>LTerm, i:Int, t:LTerm) ={ // NICE: factoring out case for variables 
    //TODO: can you generalise this into a library for dealing with de Bruijn indices?
    // i.e., clients of that library should just write a map that calls the right function for every AST construct
    // (--> provide helper for binders and variables)
   
    def walk(c: Int, t: LTerm) : LTerm = t match {
      case V(nb)        => onvar(c,V(nb))
      case L(x,ty,t1)   => L(x,ty,walk(c+1,t1))
      case A(t1,t2)     => A(walk(c,t1),walk(c,t2))
      case Tru()         => Tru
      case Fls()         => Fls
      case Unit()        => Unit
      case Lit(v)        => Lit(v)
      case If(cond,p,n)  => If(walk(c,cond),walk(c,p),walk(c,n))  
    }
    walk(0,t)
  }
   
  

}

